"""
author : DengXiuqi
date : 2018.10
email : dengxiuqi@163.com
"""

from PIL import Image, ImageEnhance
from config import *
import tensorflow as tf
from network import SRGAN_g
from utils import saveImg
import numpy as np

MEAN_VALUES = np.array([123.68, 116.779, 103.939]).reshape((1, 1, 3))


def get_test_batch():
    img = Image.open(test_img)
    img = img.resize(test_img_size)
    boxes = cut(img)
    for key, im in enumerate(boxes):
        _im = np.array(img_adjust(im, brightness, contrast, sharpness))
        _im = (_im - MEAN_VALUES) / (255 / 2)
        boxes[key] = _im
    return boxes


def img_adjust(img, brightness, contrast, sharpness):
    enh_bri = ImageEnhance.Brightness(img)
    image_brightened = enh_bri.enhance(brightness)
    enh_con = ImageEnhance.Contrast(image_brightened)
    image_contrasted = enh_con.enhance(contrast)
    enh_sha = ImageEnhance.Sharpness(image_contrasted)
    image_sharped = enh_sha.enhance(sharpness)
    return image_sharped


def cut(img):
    boxes = []
    for i in range(0, test_img_size[1], 300):
        for j in range(0, test_img_size[0], 300):
            box = (j, i, j + 300, i + 300)
            boxes.append(img.crop(box))
    return boxes


def concat(boxes, save_path=None, norm=True):
    span = test_img_size[0] // 300 + (1 if test_img_size[0] % 300 else 0)
    new_img = np.concatenate(boxes[0: 0+span], 1)
    for i in range(1, test_img_size[1] // 300 + (1 if test_img_size[1] % 300 else 0)):
        # print(i * span, (i + 1) * span)
        _img = np.concatenate(boxes[i * span: (i + 1) * span], 1)
        new_img = np.concatenate([new_img, _img])
    if norm:
        img = new_img * (255 / 2) + MEAN_VALUES
        img[img > 255] = 255
        img[img < 0] = 0
    new_img = img.astype(np.uint8)
    new_img = Image.fromarray(new_img).convert('RGB')
    if save_path:
        new_img.save(save_path)
    return new_img


def test(image_batch):
    image = tf.placeholder(dtype=tf.float32, shape=[1, img_size[0], img_size[1], 3], name='image')
    net_g = SRGAN_g(image, is_train=True, reuse=False)
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_proportion, allow_growth=True)
    config = tf.ConfigProto(gpu_options=gpu_options)

    with tf.Session(config=config) as sess:
        sess.run(tf.global_variables_initializer())
        saver = tf.train.Saver(tf.global_variables())
        saver.restore(sess, model_path)
        output = []
        for key, img in enumerate(image_batch):
            img = np.expand_dims(img, 0)
            new_img = sess.run(net_g, {image: img})
            # saveImg(new_img, key, name='output')
            output.append(new_img.reshape(300,300,-1))
    return output


if __name__ == "__main__":
    image_batch = get_test_batch()
    output = test(image_batch)
    concat(output, save_path)
